package com.ubxtech.gateway.balance;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.ubxtech.balance.dto.ServiceInstanceGreyDTO;
import com.ubxtech.balance.rule.WebSocketLoadBalancerRule;
import com.ubxtech.balance.util.ServiceInstanceUtil;
import com.ubxtech.core.context.ContextConstants;
import com.ubxtech.core.utils.GrayVersionUtil;
import com.ubxtech.core.utils.SpringUtils;
import com.ubxtech.gateway.properties.CustomGatewayProperties;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.cloud.gateway.support.NotFoundException;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.util.MultiValueMap;

import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 版本号灰度路由,载均衡策略的实现
 * @author Victor.Xiao
 * @since 2023-08-22 14:54
*/
@Slf4j
public class VersionGrayLoadBalancerRule implements GrayLoadBalancerRule {
	private DiscoveryClient discoveryClient;
	private CustomGatewayProperties customGatewayProperties;

	/** 用于轮训算法 */
	private AtomicInteger nextServerAtomic;

	private WebSocketLoadBalancerRule webSocketLoadBalancerRule;

	public VersionGrayLoadBalancerRule(DiscoveryClient discoveryClient, CustomGatewayProperties customGatewayProperties) {
		this.discoveryClient = discoveryClient;
		this.customGatewayProperties = customGatewayProperties;
		nextServerAtomic = new AtomicInteger(0);
	}

	/**
	 * 根据serviceId 筛选可用服务
	 *
	 * @param serviceId 服务ID
	 * @param request   当前请求
	 * @return
	 */
	@Override
	public ServiceInstance choose(String serviceId, ServerHttpRequest request) {
		List<ServiceInstance> instances = discoveryClient.getInstances(serviceId);

		//注册中心无实例 抛出异常
		if (CollUtil.isEmpty(instances)) {
			log.warn("No instance available for {}", serviceId);
			throw new NotFoundException("No instance available for " + serviceId);
		}
		// 获取请求version，无则随机返回可用实例
		String reqVersion = request.getHeaders().getFirst(ContextConstants.HEADER_GRAY_VERSION_NAME);
		//分析出有版本号的实例
		ServiceInstanceGreyDTO serviceInstanceGreyDTO = ServiceInstanceUtil.analysisGreyInstance(instances, reqVersion);
		//匹配的实例
		List<ServiceInstance> instancesFilter = serviceInstanceGreyDTO.getInstancesFilter();
		//version为空或者为all的实例,目的是如果没有带版本号,其中有服务带了版本号,则不调用
		List<ServiceInstance> emptyVersionInstancesFilter = serviceInstanceGreyDTO.getEmptyVersionInstancesFilter();

		String upgrade = request.getHeaders().getFirst(ContextConstants.HEADER_UPGRADE);
		//如果是websocket,单独处理转发规则,如果要支持接口也走websocket的转发规则,则要放开这里,额外根据服务判断或者加入其他逻辑判断
		//目前接口不对外提供,只处理feign端的转发规则,根据serviceId前缀和url判断,避免整个服务的接口都是同一个转发规则
		if (ContextConstants.HEADER_UPGRADE_VALUE.equalsIgnoreCase(upgrade)) {
			String userId = request.getHeaders().getFirst(ContextConstants.USER_ID);
			String token = request.getHeaders().getFirst(ContextConstants.TOKEN);
			//websocket可能带sn,以sn为准
			MultiValueMap<String, String> queryPrams =  request.getQueryParams();
			String sn = queryPrams.getFirst(ContextConstants.PARAM_SN);
			if (StrUtil.isNotBlank(sn)) {
				userId = sn;
			}
			if (webSocketLoadBalancerRule == null ) {
				if (SpringUtils.hasBean("webSocketLoadBalancerRule")) {
					webSocketLoadBalancerRule = SpringUtils.getBean(WebSocketLoadBalancerRule.class);
				}
			}
			if (webSocketLoadBalancerRule != null) {
				ServiceInstance serviceInstance = webSocketLoadBalancerRule.choose(instancesFilter, emptyVersionInstancesFilter, instances, userId, token);
				if (serviceInstance != null) {
					return serviceInstance;
				}
			}
		}

		//优先匹配到的,没有匹配到再是没有配置版本的,然后再是所有
		return GrayVersionUtil.getInstance(instancesFilter.size() == 0 ? (emptyVersionInstancesFilter.size() == 0 ? instances : emptyVersionInstancesFilter) : instancesFilter, customGatewayProperties.getBalanceType(), nextServerAtomic);
//		return GrayVersionUtil.getInstance(instancesFilter.size() == 0 ? instances : instancesFilter, customGatewayProperties.getBalanceType(), nextServerAtomic);
	}

}
